/*
 * subgradient_parallel.cc
 *
 *  Created on: Sep 14, 2011
 *      Author: tqlong
 */
#include <boost/foreach.hpp>
#include "subgradient_parallel.h"
#include "lp_lib.h"

SubGradGM::SubGradGM(const SimpleGM& gm)
{
  cardinality_ = gm.cardinality_;
  node_potentials_ = gm.node_potentials_;
  edge_potentials_ = gm.edge_potentials_;
  adj_edges_ = gm.adj_edges_;
  node_marginals_ = gm.node_marginals_;
  edge_marginals_ = gm.edge_marginals_;
  old2new_ = gm.old2new_;
  new2old_ = gm.new2old_;
}

SubGradGM::assignment_type SubGradGM::subgradient_lp_map()
{
  init_variables();

  eta_ = 1.0;
  double minval = std::numeric_limits<double>::infinity();
  int no_change = 0;
  n_potential_type minmu;
  val_ = std::numeric_limits<double>::infinity();

  solve_slave_problems(var_);  // fix nu_ and nue_ first
  for (int iter = 0; iter < 500 && no_change < 100 && eta_ > 1e-6; iter++)
  {
    eta_ = 1.0;
    val_ = objective(var_);
    if (minval > val_)
    {
      minval = val_;
      minmu = var_.mu_;
      no_change = 0;
    }
    else
      no_change++;


    tmp_.lambda_ = var_.lambda_;
    solve_slave_problems(tmp_);    // to find gradient  \nabla_\theta = \nu
    double s;
    update_mu(tmp_, s);            // to find search direction   d(i,e) = (\theta(i,e) - \mu(i))

    //do_line_search(val_, s);          // to find step size \eta

    //subgradient without line search
    eta_ = 1.0 / sqrt(iter+1);
    BOOST_FOREACH(pair_nodeedge_t p, var_.lambda_)
    {
      p.second -= eta_ * d_[p.first];
    }
    solve_slave_problems(var_);

    std::cout << " ---------  iter=" << iter << " obj="  << val_
              << " no_change=" << no_change << " eta=" << eta_ << std::endl;
//    BOOST_FOREACH(const_pair_node_t p, mu_)
//    {
//      std::cout << "mu(" << p.first << ")=" << std::endl << p.second;
//    }
  }

  var_.mu_ = minmu;
  assignment_type a;
  bool non_integer_solution = false;
  BOOST_FOREACH(const_pair_node_t p, var_.mu_)
  {
    Node i = p.first;
    const arma::vec& v = p.second; // the max marginal
    int val = 0;
    for (int vi = 1; vi < cardinality_[i]; vi++)
      if (v[vi] > v[val]) val = vi;
    a[i] = val;
    if (v[val] < 1.0-1e-5)
    {
      non_integer_solution = true;
    }
  }
  std::cout << ( non_integer_solution ? "not an integer solution" : "integer optimality reached" )
            << std::endl;
  return a;
}

// Protected members
void SubGradGM::solve_slave_problems(optim_var_type& var)
{
  BOOST_FOREACH(pair_edge_t p, var.nue_)
  {
    Edge e = p.first;
    arma::mat& ev = p.second;  // nue_[e]
    Node i = e.i_, j = e.j_;
    arma::vec wi, wj;
    const arma::mat& phi = edge_potentials_[e];
    build_slave_problem(var, i, j, e, wi, wj);

//    std::cout << "wi(" << i << ")=" << std::endl << wi
//              << "wj(" << j << ")=" << std::endl << wj;

    arma::vec nui, nuj;
    arma::mat nue;
    solve_slave_problem(wi, wj, phi, nui, nuj, nue);

//    std::cout << "nui(" << i << ", (" << e.i_ << "," << e.j_ << "))=" << std::endl << nui
//              << "nuj(" << j << ", (" << e.i_ << "," << e.j_ << "))=" << std::endl << nuj;

    var.nu_[NodeEdge(i, e)] = nui;
    var.nu_[NodeEdge(j, e)] = nuj;
    ev = nue;
  }

  BOOST_FOREACH(pair_node_t p, var.mu_)
  {
    Node i = p.first;
    arma::vec& mu_i = p.second;
    mu_i.zeros();
    BOOST_FOREACH(Edge e, adj_edges_[i])
    {
      NodeEdge ne(i, e);
      mu_i += var.nu_[ne];
    }
    mu_i /= adj_edges_[i].size();
  }
}

void SubGradGM::init_variables()
{
  var_.mu_ = node_potentials_;
  var_.nue_ = edge_potentials_;

  BOOST_FOREACH(pair_node_t p, var_.mu_)
  {
    p.second.zeros();
  }

  BOOST_FOREACH(pair_edge_t p, var_.nue_)
  {
    p.second.zeros();
  }

  BOOST_FOREACH(const_pair_adje_t p, adj_edges_)
  {
    const Node& i = p.first;
    BOOST_FOREACH(const Edge& e, p.second)
    {
      NodeEdge ne(i, e);
      var_.nu_[ne] = arma::zeros<arma::vec>(cardinality_[i]);
      var_.lambda_[ne] = arma::zeros<arma::vec>(cardinality_[i]);
    }
  }
  tmp_ = var_;
}

void SubGradGM::build_slave_problem(optim_var_type& var, Node i, Node j, Edge e, arma::vec& wi, arma::vec& wj)
{
  wi = node_potentials_[i] / adj_edges_[i].size() + var.lambda_[NodeEdge(i,e)];
  wj = node_potentials_[j] / adj_edges_[j].size() + var.lambda_[NodeEdge(j,e)];
}

typedef std::pair<int, int> grid_t;
typedef std::map<int, int> idx2idx_t;
typedef std::map<grid_t, int> grid2idx_t;

lprec* build_lp_model(const arma::vec& wi, const arma::vec& wj, const arma::mat& phi,
    idx2idx_t& nui_map, idx2idx_t& nuj_map, grid2idx_t& nue_map, int& ncol)
{
# if defined RETURN
#  undef RETURN
# endif
# define RETURN(lp) { delete row; delete colno; return lp; }

//  std::cout << "wi=" << std::endl << wi;
//  std::cout << "wj=" << std::endl << wj;
//  std::cout << "phi=" << std::endl << phi;

  int ni = wi.n_elem, nj = wj.n_elem;
  ncol = ni + nj + phi.n_elem;

  lprec* lp = make_lp(0, ncol);
  int *colno = new int [ncol];
  REAL *row = new REAL [ncol];

  // variable order: nui, nuj, nue (by columns)
  int k = 1;
  char buf[100];
  for (int i = 0; i < ni; i++, k++)
  {
    nui_map[i] = k;
    sprintf(buf, "nui_%d", i);
    set_col_name(lp, k, buf);
  }
  for (int j = 0; j < nj; j++, k++)
  {
    nuj_map[j] = k;
    sprintf(buf, "nuj_%d", j);
    set_col_name(lp, k, buf);
  }
  for (int j = 0; j < nj; j++)
    for (int i = 0; i < ni; i++, k++)
    {
      nue_map[grid_t(i,j)] = k;
      sprintf(buf, "nue_(%d,%d)", i, j);
      set_col_name(lp, k, buf);
    }

  // objective
  set_add_rowmode(lp, FALSE);
  k = 0;
  for (int i = 0; i < ni; i++)
  {
    colno[k] = nui_map[i];
    row[k++] = wi[i];
  }
  for (int j = 0; j < nj; j++)
  {
    colno[k] = nuj_map[j];
    row[k++] = wj[j];
  }
  for (int j = 0; j < nj; j++)
    for (int i = 0; i < ni; i++)
    {
      colno[k] = nue_map[grid_t(i,j)];
      row[k++] = phi(i,j);
    }
  if(!set_obj_fnex(lp, k, row, colno))
    RETURN(NULL);

  // box constraints >= 0 <= 1
  for (int i = 0; i < ni; i++)
  {
    if (!set_bounds(lp, nui_map[i], 0.0, 1.0))
      RETURN(NULL);
  }
  for (int j = 0; j < nj; j++)
  {
    if (!set_bounds(lp, nuj_map[j], 0.0, 1.0))
      RETURN(NULL);
  }
  for (int i = 0; i < ni; i++)
    for (int j = 0; j < nj; j++)
    {
      if (!set_bounds(lp, nue_map[grid_t(i,j)], 0.0, 1.0))
        RETURN(NULL);
    }

  // constraints
  set_add_rowmode(lp, TRUE);

  // < w_i , 1 > = 1
  k = 0;
  for (int i = 0; i < ni; i++)
  {
    colno[k] = nui_map[i];
    row[k++] = 1;
  }
  if(!add_constraintex(lp, k, row, colno, EQ, 1.0))
    RETURN(NULL);

  // < w_j , 1 > = 1
  k = 0;
  for (int j = 0; j < nj; j++)
  {
    colno[k] = nuj_map[j];
    row[k++] = 1;
  }
  if(!add_constraintex(lp, k, row, colno, EQ, 1.0))
    RETURN(NULL);

  // w_i = phi * 1
  for (int i = 0; i < ni; i++)
  {
    k = 0;
    colno[k] = nui_map[i];
    row[k++] = -1;
    for (int j = 0; j < nj; j++)
    {
      colno[k] = nue_map[grid_t(i,j)];
      row[k++] = 1;
    }
    if(!add_constraintex(lp, k, row, colno, EQ, 0.0))
      RETURN(NULL);
  }

  // w_j' = 1' * phi
  for (int j = 0; j < nj; j++)
  {
    k = 0;
    colno[k] = nuj_map[j];
    row[k++] = -1;
    for (int i = 0; i < ni; i++)
    {
      colno[k] = nue_map[grid_t(i,j)];
      row[k++] = 1;
    }
    if(!add_constraintex(lp, k, row, colno, EQ, 0.0))
      RETURN(NULL);
  }

  set_add_rowmode(lp, FALSE);
  set_maxim(lp);
  RETURN(lp);
}

void SubGradGM::solve_slave_problem(const arma::vec& wi, const arma::vec& wj, const arma::mat& phi,
                                    arma::vec& nui, arma::vec& nuj, arma::mat& nue)
{
# if defined ERROR
#  undef ERROR
# endif
# define ERROR() { fprintf(stderr, "Error\n"); exit(1); }

  idx2idx_t nui_map, nuj_map;
  grid2idx_t nue_map;
  int ncol;
  lprec *lp  = build_lp_model(wi, wj, phi, nui_map, nuj_map, nue_map, ncol);
  if (!lp)
  {
    ERROR();
  }

//  print_lp(lp);
  set_verbose(lp, IMPORTANT);
  int ret = solve(lp);

  if (ret != OPTIMAL)
  {
    ERROR();
  }
  else // collect optimized variables;
  {
    REAL *row = new REAL[ncol];
    get_variables(lp, row);
    nui = wi;
    nuj = wj;
    nue = phi;
    for (unsigned int i = 0; i < nui.n_elem; i++) nui[i] = row[nui_map[i]-1]; //-1 as lpsolve use index
    for (unsigned int j = 0; j < nuj.n_elem; j++) nuj[j] = row[nuj_map[j]-1]; // starting from 1
    for (unsigned int i = 0; i < nui.n_elem; i++)
      for (unsigned int j = 0; j < nuj.n_elem; j++)
        nue.at(i,j) = row[nue_map[grid_t(i,j)]-1];

//    for(int j = 0; j < ncol; j++)
//      std::cout << get_col_name(lp, j + 1) << " = " << row[j] << " ";
//    std::cout << std::endl;
    delete row;
  }
  delete_lp(lp);
}

// mu_i = average_e(nu(i, e))
void SubGradGM::update_mu(optim_var_type& var, double& s)
{
  d_ = var.nu_;
  s = 0;
  typedef duplicate_variable_type::value_type& pair_nodeedge_t;
  BOOST_FOREACH(pair_nodeedge_t p, d_)
  {
    p.second -= var.mu_[p.first.first];      // d(i,e) = \lambda(i,e) - \mu(i)
    s += arma::dot(p.second, p.second);      // the decreasing rate along direction d
  }
}

// lambda(i,e) <-- lambda(i,e) - eta_t (nu(i,e) - mu(i))
void SubGradGM::update_lambda()
{
}

double SubGradGM::objective(optim_var_type& var)
{
  double s = 0;
  BOOST_FOREACH(pair_edge_t p, var.nue_)
  {
    Edge e = p.first;
    Node i = e.i_, j = e.j_;
    s += arma::dot(p.second, edge_potentials_[e]);
    s += arma::dot(node_potentials_[i] / adj_edges_[i].size() + var.lambda_[NodeEdge(i,e)], var.nu_[NodeEdge(i,e)]);
    s += arma::dot(node_potentials_[j] / adj_edges_[j].size() + var.lambda_[NodeEdge(j,e)], var.nu_[NodeEdge(j,e)]);
  }
  return s;
}

void SubGradGM::do_line_search(double f0, double s)
{
  double min_eta = 0;
  double min_val = f0;
  while (eta_ > 1e-6)
  {
    BOOST_FOREACH(pair_nodeedge_t p, tmp_.lambda_)
    {
      p.second = var_.lambda_[p.first] - eta_ * d_[p.first];
    }
    solve_slave_problems(tmp_);
    val_ = objective(tmp_);
    if (val_ < min_val)
    {
      min_val = val_;
      min_eta = eta_;
    }
    std::cout << "  eta=" << eta_ << " val=" << val_  << " f0=" << f0 << " need=" << f0 - 0.5*eta_*s << std::endl;
    if (val_ <= f0 - 0.5*eta_*s - 1e-6)
    {
      var_.lambda_ = tmp_.lambda_;
      var_.nu_ = tmp_.nu_;
      var_.nue_ = tmp_.nue_;
      var_.mu_ = tmp_.mu_;
      std::cout << "GOOD" << std::endl;
      break;
    }

    eta_ /= 2.0;
    if (eta_ < 1e-6)
    {
      eta_ = min_eta;
      val_ = min_val;
      BOOST_FOREACH(pair_nodeedge_t p, var_.lambda_)
      {
        p.second -= eta_ * d_[p.first];
      }
      solve_slave_problems(var_);
      std::cout << "BAD BAD eta=" << eta_ << std::endl;
      break;
    }
  }
}
